"""
没有声发射数据，所以这部分还是振动的逻辑
"""
from torch.utils.data import Dataset, DataLoader
import numpy as np
from torch.utils.data import WeightedRandomSampler
import torch
import pywt


class VibrationDataset(Dataset):
    def __init__(self, data_path):
        self.data = np.load(data_path)
        self.vibration = self.data[:, :-1]
        self.label = self.data[:, -1]
        self.indices = np.arange(self.vibration.shape[0], dtype=int)
        self.nc = len(np.unique(self.label))

    def __len__(self):
        return len(self.vibration)

    def __getitem__(self, index):

        idx = self.indices[index]
        # Add noise to the data
        # The SNR is randomly selected from -5 to 5 dB
        # if np.random.random() < 0.2:
        #     vib_power = (np.abs(self.vibration[idx]) ** 2).mean()
        #     snr = np.random.randint(0, 5)
        #     noise_power = vib_power / (10 ** (snr / 10))
        #     noise = self.gen_noise(np.sqrt(noise_power), self.vibration[idx].shape[0])
        #     self.vibration[idx] += noise
        
        # Add a slide-average filter to the vibration data
        # elif np.random.random() < 0.2:
        #     self.vibration[idx] = np.convolve(self.vibration[idx], np.ones(8) / 8, mode='same')
        
        # normalize the vibration data
        self.vibration[idx] = (self.vibration[idx] - self.vibration[idx].mean()) / self.vibration[idx].std()
        return torch.from_numpy(self.vibration[idx]).type(torch.float32), self.label[idx]
    
    @staticmethod
    def gen_noise(std, size, mean=0):
        return np.random.normal(mean, std, size)


if __name__ == '__main__':
    my_dataset = VibrationDataset('SAFDNN/test_data.npy')

    # 计算每个类别的权重
    class_counts = np.bincount(my_dataset.label.astype(int))
    class_weights = 1. / class_counts
    sample_weights = class_weights[my_dataset.label.astype(int)]
    
    # 创建 WeightedRandomSampler
    sampler = WeightedRandomSampler(weights=sample_weights, num_samples=len(sample_weights), replacement=False)
    
    # 创建 DataLoader
    dataloader = DataLoader(my_dataset, batch_size=32)
    for epoch in range(4):
        for idx, (data, target) in enumerate(dataloader):
            print(data.shape, target.shape)
            if idx == 3:
                break
        print("Epoch", epoch)
        print(dataloader.dataset.indices)
        dataloader.dataset.indices = np.ones_like(dataloader.dataset.indices)

